/*
 *  jOTR - Java Off-The-Record Messaging Library
 *  Copyright (C) 2007-2008 Markus Karnik
 *
 *  This program is free software; you can redistribute it and/or modify it 
 *  under the terms of the GNU General Public License as published by the
 *  Free Software Foundation; either version 3 of the License, 
 *  or (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful, but
 *  WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
 *  or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License 
 *  for more details.
 *
 *  You should have received a copy of the GNU General Public License along 
 *  with this program; if not, see <http://www.gnu.org/licenses/>.
 *  
 */
package de.karnik.libjotr;

import java.math.BigInteger;
import java.security.InvalidAlgorithmParameterException;
import java.security.NoSuchAlgorithmException;

import de.karnik.helper.ArrayFunctions;
import de.karnik.helper.HashFunctions;
import de.karnik.helper.LogFunctions;
import de.karnik.libjotr.helper.JotrFunctions;

/**
 * This class serves with Diffie-Hellmann key computation and some helper functions.
 * 
 * @author Markus Karnik - markus.karnik@my-design.org
 * @version 1.0
 * @since 1.0
 */
public class JotrDH implements JotrConstants {
	
	/**
	 * The Diffie-Hellman modulus value.
	 */
	private BigInteger modulus = null;
	/**
	 * The Diffie-Hellman generator value.
	 */
	private BigInteger generator = null;
	
	/**
	 * The bit-length of the exponent.
	 */
	private int length = 0;
    
    /**
     * The current jotr auth keys object. (user for ake phase)
     */
    private JotrAuthKeys authKeys = null;
    
    /**
     * The shared secret.
     */
    private JotrMpi sharedSecretKey = null;
    
    /**
     * The old mac keys to reveal.
     */
    private byte[] revealedMacKeys = null;
    
	/**
	 * The previous local dh key pair.
	 */
	private JotrDHKeyPair prevLocalDHKeyPair = null;
	/**
	 * The previous remote dh key.
	 */
	private JotrMpi prevRemotePublicKey = null;
	/**
	 * The previous remote dh key id.
	 */
	private int prevRemotePublicKeyId = 0;
	
	/**
	 * The current local dh key pair.
	 */
	private JotrDHKeyPair curLocalDHKeyPair = null;	
	/**
	 * The current remote dh key.
	 */
	private JotrMpi curRemotePublicKey = null;
	/**
	 * The current remote dh key id.
	 */
	private int curRemotePublicKeyId = 0;
	
    /**
     * Sesskeys[i][j] are the session keys derived from DH key[our_keyid-i] and mpi Y[their_keyid-j]!
     */
    private JotrSessionKeys[][] sessKeys = new JotrSessionKeys[ 2 ][ 2 ];
	
	/**
	 * Standard constructor. Hidden.
	 */
	private JotrDH() {}

	/**
	 * Constructor to create a JotrDH object with the given generator, modulus value, exponent bit count and ip of the modpow server.
	 * 
	 * This a special constructor for  the workaround caused by performance issues.
	 * 
	 * @param generator The generator for the dh computations.
	 * @param modulus The modulus for the dh computations.
	 * @param bitCount The length of the exponent in bit.
	 * @param ip The ip of the modpow server.
	 * @throws InvalidAlgorithmParameterException
	 * @throws NoSuchAlgorithmException
	 * 
	 * @see de.karnik.libjotr.JotrConstants#DH1536_GENERATOR_S
	 * @see de.karnik.libjotr.JotrConstants#DH1536_MODULUS_S
	 */
	public JotrDH( String generator, String modulus, int bitCount, String ip ) throws InvalidAlgorithmParameterException, NoSuchAlgorithmException {
		this( generator, modulus, bitCount );
	}
	
	/**
	 * Constructor to create a JotrDH object with the given generator, modulus value and exponent bit count.
	 * 
	 * @param generator The generator for the dh computations.
	 * @param modulus The modulus for the dh computations.
	 * @param bitCount The length of the exponent in bit.
	 * @throws InvalidAlgorithmParameterException
	 * @throws NoSuchAlgorithmException
	 * 
	 * @see de.karnik.libjotr.JotrConstants#DH1536_GENERATOR_S
	 * @see de.karnik.libjotr.JotrConstants#DH1536_MODULUS_S
	 */
	public JotrDH( String generator, String modulus, int bitCount ) throws InvalidAlgorithmParameterException, NoSuchAlgorithmException {
		this();
		this.generator = new BigInteger ( generator, 10 );
		this.modulus = new BigInteger( modulus, 16 );
		this.length = bitCount;
		
		this.curLocalDHKeyPair = new JotrDHKeyPair( null, null );
		this.prevLocalDHKeyPair = new JotrDHKeyPair( null, null );
	}

	// #############################################################################
	// Init-Functions
	// #############################################################################
	
	/**
	 * Generates a length-bit random BigInteger value. Length is the length of the secret key. Default is 320!
	 * 
	 * @return A length-bit random BigInteger value.
	 */
	private BigInteger generateSecretKey() {
		// process exponent
		byte[] mag;
	    BigInteger x;
	    while ( true ) {
	    	mag = ArrayFunctions.generateRandomByteArray( ( length + 7 ) / 8 );
	    	
	    	//LogFunctions.log( "JotrDH.computeDiffieHellmanKeyPair" , "compute x", StringFunctions.bytesToHex( mag ) );
	    	
	    	x = new BigInteger( 1, mag );
	        if ( x.bitLength() == length 
	        		&& x.compareTo( BigInteger.ONE ) > 0 
	        		&& x.compareTo( modulus.subtract( BigInteger.ONE) ) < 0 )
	        	break;
	    }
	    
	    return x;
	}
	
	/**
	 * Computes d-h shared secret ( s = (gx)y ) ( this is needed to decrypt the signature ) and 
	 * two AES keys c, c' and four MAC keys m1, m1', m2, m2' by hashing s in various ways.
	 */
	public void computeAKESharedSecretKeys() {
		
		// compute and get shared secret
		computeSharedSecret( 	prevLocalDHKeyPair.getPrivateKey().getValueAsBigInteger(), 
								curRemotePublicKey.getValueAsBigInteger() );
		
		byte[] cKey =  JotrFunctions.computeAKEOtrKeys( KEY_C, sharedSecretKey );
		byte[] ciKey = JotrFunctions.computeAKEOtrKeys( KEY_C_I, sharedSecretKey );
		byte[] m1Key = JotrFunctions.computeAKEOtrKeys( KEY_M1, sharedSecretKey );
		byte[] m2Key =  JotrFunctions.computeAKEOtrKeys( KEY_M2, sharedSecretKey );
		byte[] m1iKey =  JotrFunctions.computeAKEOtrKeys( KEY_M1_I, sharedSecretKey );
		byte[] m2iKey = JotrFunctions.computeAKEOtrKeys( KEY_M2_I, sharedSecretKey );
		
		authKeys = new JotrAuthKeys( cKey, ciKey, m1Key, m2Key, m1iKey, m2iKey); 
	}
	
	/**
	 * Computes d-h shared secret ( s = (gx)y ) and two AES keys sending, receiving and tow 
	 * MAC keys sending receiving by hashing s in various ways.
	 * 
	 * @param dhKeys
	 * @param remotePubKey
	 * @return JotrSessionKeys
	 */
	public JotrSessionKeys computeSessionKeys( JotrDHKeyPair dhKeys, BigInteger remotePubKey ) {
				
		BigInteger localPrivateKey = dhKeys.getPrivateKey().getValueAsBigInteger();
		BigInteger localPublicKey = dhKeys.getPublicKey().getValueAsBigInteger();
		
		// compute new shared secret
		computeSharedSecret( localPrivateKey, remotePubKey );
		
		byte sendbyte = ( byte )0x01;
		byte recvbyte = ( byte )0x02;
		
		// determine if "high" or "low"
		if( localPublicKey.compareTo( remotePubKey ) == -1 ) {
			// low
			sendbyte = ( byte )0x02;
			recvbyte = ( byte )0x01;
		}
		
		// compute sending key and mac key
	    byte[] sendingAesKey = JotrFunctions.computeSendRecvAesKey( sendbyte, sharedSecretKey );
	    byte[] sendingMacKey = HashFunctions.sha1( sendingAesKey );

		// compute receiving key and mac key
	    byte[] receivingAesKey = JotrFunctions.computeSendRecvAesKey( recvbyte, sharedSecretKey );
	    byte[] receivingMacKey = HashFunctions.sha1( receivingAesKey );

	    // create and return new session key set
	    return new JotrSessionKeys( sendingAesKey, receivingAesKey, sendingMacKey, receivingMacKey );
	}
	
	/**
	 * Searches for used mac keys and stores the result to revealedMacKeys.
	 * 
	 * @param keys1 The first JotrSessionKeys object to search for used mac keys.
	 * @param keys2 The second JotrSessionKeys object to search for used mac keys.
	 * 
	 * @see #revealedMacKeys
	 */
	private void reveal( JotrSessionKeys keys1, JotrSessionKeys keys2 ) {
				
		if( keys1 != null && keys1.isReceivingUsed() )
			revealedMacKeys = ArrayFunctions.merge( revealedMacKeys, keys1.getReceivingMacKey() );
			
		if( keys1 != null && keys1.isSendingUsed() )
			revealedMacKeys = ArrayFunctions.merge( revealedMacKeys, keys1.getSendingMacKey() );

		if( keys2 != null && keys2.isReceivingUsed() )
			revealedMacKeys = ArrayFunctions.merge( revealedMacKeys, keys2.getReceivingMacKey() );
			
		if( keys2 != null && keys2.isSendingUsed() )
			revealedMacKeys = ArrayFunctions.merge( revealedMacKeys, keys2.getSendingMacKey() );	
	}
	
	/**
	 * Sets the old remote key pair to the current one and computes a new D-H Key pair. 
	 * 
	 * @param newRemotePubKey
	 */
	public void rotateRemotePublicKey( byte[] newRemotePubKey ) {
		
		// save old one
		prevRemotePublicKeyId = curRemotePublicKeyId;
		prevRemotePublicKey = curRemotePublicKey;
		
		// reveal old keys
		reveal( sessKeys[ 0 ][ 1 ], sessKeys[ 1 ][ 1 ] );
		
		// save session keys
		sessKeys[ 0 ][ 1 ] = sessKeys[ 0 ][ 0 ];
		sessKeys[ 1 ][ 1 ] = sessKeys[ 1 ][ 0 ];
		
		// set new remote key
		curRemotePublicKey = new JotrMpi( newRemotePubKey );
		curRemotePublicKeyId++;
		
		// create the session keys
	    sessKeys[ 0 ][ 0 ] = computeSessionKeys( curLocalDHKeyPair, curRemotePublicKey.getValueAsBigInteger() );
	    sessKeys[ 1 ][ 0 ] = computeSessionKeys( prevLocalDHKeyPair, curRemotePublicKey.getValueAsBigInteger() );
		
	}
	
	/**
	 * Sets the old local key pair to the current one and computes session keys for the new one.
	 */
	public void rotateLocaleKeyPairs() {
		
		JotrMpi x;

		// save old ones
		prevLocalDHKeyPair.setPrivateKey( curLocalDHKeyPair.getPrivateKey() );
		prevLocalDHKeyPair.setPublicKey( curLocalDHKeyPair.getPublicKey() );
		prevLocalDHKeyPair.setKeyID( curLocalDHKeyPair.getKeyID() );
		
		LogFunctions.log( "JotrDH.rotateLocaleKeyPairs", " prevLocalKeyId", prevLocalDHKeyPair.getKeyID() );
		LogFunctions.log( "JotrDH.rotateLocaleKeyPairs", " curLocalKeyId", curLocalDHKeyPair.getKeyID() );
		
		
		// reveal old keys
		reveal( sessKeys[ 1 ][ 0 ], sessKeys[ 1 ][ 1 ] );
		
		// save session keys
	    this.sessKeys[ 1 ][ 0 ] = sessKeys[ 0 ][ 0 ];
	    this.sessKeys[ 1 ][ 1 ] = sessKeys[ 0 ][ 1 ];
			
		// generate new key pair
		x = new JotrMpi( generateSecretKey() );
		curLocalDHKeyPair = new JotrDHKeyPair( new JotrMpi( generator.modPow( x.getValueAsBigInteger(), modulus ) ), x );
			
		// increment key id 
		curLocalDHKeyPair.setKeyID( prevLocalDHKeyPair.getKeyID() + 1 );
		
		// create new sessionkeys
	    if ( curRemotePublicKey != null )
	    	sessKeys[ 0 ][ 0 ] = computeSessionKeys( curLocalDHKeyPair, curRemotePublicKey.getValueAsBigInteger() );
	    else 
	    	sessKeys[ 0 ][ 0 ] = null;
	    
	    if( prevRemotePublicKey != null )
	    	sessKeys[ 0 ][ 1 ] = computeSessionKeys( curLocalDHKeyPair, prevRemotePublicKey.getValueAsBigInteger() );
	    else 
	    	sessKeys[ 0 ][ 1 ] = null;
	}
	
	public void computeInitialSessionKeys() {
		
		// create the session keys
		
	    if ( curRemotePublicKey != null ) {
	    	sessKeys[ 0 ][ 0 ] = computeSessionKeys( curLocalDHKeyPair, curRemotePublicKey.getValueAsBigInteger() );
	    	sessKeys[ 1 ][ 0 ] = computeSessionKeys( prevLocalDHKeyPair, curRemotePublicKey.getValueAsBigInteger() );
	    } else { 
	    	sessKeys[ 0 ][ 0 ] = null;
	    	sessKeys[ 1 ][ 0 ] = null;
	    }
	    
	    if( prevRemotePublicKey != null ) {
	    	sessKeys[ 0 ][ 1 ] = computeSessionKeys( curLocalDHKeyPair, prevRemotePublicKey.getValueAsBigInteger() );
	    	sessKeys[ 1 ][ 1 ] = computeSessionKeys( prevLocalDHKeyPair, prevRemotePublicKey.getValueAsBigInteger() );
	    } else { 
	    	sessKeys[ 0 ][ 1 ] = null;
	    	sessKeys[ 1 ][ 1 ] = null;
	    }
	}
	
	/**
	 * Computes the D-H shared secret from the given keys.
	 * 
	 * @param privateKey The private key to calculate the shared secret.
	 * @param publicKey The public key to calculate the shared secret.
	 */
	public void computeSharedSecret( BigInteger privateKey, BigInteger publicKey ) {
		
		if( privateKey != null && publicKey != null ) {
				sharedSecretKey = new JotrMpi( publicKey.modPow( privateKey, modulus ) );
		}
	}
	
	/**
	 * Computes two fresh D-H key pairs.
	 */
	public void computeFreshDiffieHellmanKeyPairs() {

		JotrMpi x;
	
		// generate new key pair
		x = new JotrMpi( generateSecretKey() );
		prevLocalDHKeyPair = new JotrDHKeyPair( new JotrMpi( generator.modPow( x.getValueAsBigInteger(), modulus ) ), x );
		prevLocalDHKeyPair.setKeyID( 1 );

		// generate new key pair
		x = new JotrMpi( generateSecretKey() );
		curLocalDHKeyPair = new JotrDHKeyPair( new JotrMpi( generator.modPow( x.getValueAsBigInteger(), modulus ) ), x );
		curLocalDHKeyPair.setKeyID( 2 );
	}

	// #############################################################################
	// Setter and Getter
	// #############################################################################
	
	/**
	 * Returns the dh modulus value.
	 * 
	 * @return The dh modulus value.
	 */
	public BigInteger getModulus() {
		return modulus;
	}

	/**
	 * Returns the dh generator value.
	 * 
	 * @return The dh generator value.
	 */
	public BigInteger getGenerator() {
		return generator;
	}

	/**
	 * Returns the length in bit of the dh exponent.
	 * 
	 * @return The length in bit of the dh exponent.
	 */
	public int getLength() {
		return length;
	}
	
	/**
	 * Returns the current remote dh public key.
	 * 
	 * @return The current remote dh public key.
	 */
	public JotrMpi getCurrentRemotePublicKey() {
		return curRemotePublicKey;
	}
	
	/**
	 * Returns the previous remote dh public key.
	 * 
	 * @return The previous remote dh public key.
	 */
	public JotrMpi getPreviousRemotePublicKey() {
		return prevRemotePublicKey;
	}
	
	/**
	 * Returns the current remote dh public key as byte[].
	 * 
	 * @return The current remote dh public key as byte[].
	 */
	public byte[] getCurrentRemotePublicKeyAsBytes() {
		return curRemotePublicKey.getValueAsByteArray();
	}

	/**
	 * Returns the previous remote dh public key as byte[].
	 * 
	 * @return The previous remote dh public key as byte[].
	 */
	public byte[] getPreviousRemotePublicKeyAsBytes() {
		return prevRemotePublicKey.getValueAsByteArray();
	}
	
	/**
	 * Sets the current remote dh public key.
	 * 
	 * @param remotePublicKeyBytes The current remote dh public key.
	 */
	public void setCurrentRemotePublicKey( byte[] remotePublicKeyBytes ) {
		curRemotePublicKey = new JotrMpi( remotePublicKeyBytes );
	}
	
	/**
	 * Sets the previous remote dh public key.
	 * 
	 * @param remotePublicKeyBytes The previous remote dh public key.
	 */	
	public void setPreviousRemotePublicKey( byte[] remotePublicKeyBytes ) {
	    prevRemotePublicKey = new JotrMpi( remotePublicKeyBytes );
	}
	
	/**
	 * Sets the previous remote dh public key and keyid.
	 * 
	 * @param keyID The previous remote dh public keyid.
	 * @param remotePublicKey The previous remote dh public key.
	 */
	public void setPreviousRemotePublicKey( int keyID, JotrMpi remotePublicKey ) {
		prevRemotePublicKey = remotePublicKey;
	}

	/**
	 * Returns the shared secret.
	 *  
	 * @return The shared secret.
	 */
	public JotrMpi getSharedSecretKey() {
		return sharedSecretKey;
	}

	/**
	 * Sets the shared secret.
	 * 
	 * @param secretKey The shared secret to set.
	 */
	public void setSharedSecretKey( JotrMpi secretKey ) {
		this.sharedSecretKey = secretKey;
	}

	/**
	 * Returns the previous remote dh public key id.
	 * 
	 * @return The previous remote dh public key id.
	 */
	public int getPreviousRemotePublicKeyId() {
		return prevRemotePublicKeyId;
	}

	/**
	 * Sets the previous remote dh public key id.
	 * 
	 * @param prevRemotePublicKeyId The previous remote dh public key id to set.
	 */
	public void setPreviousRemotePublicKeyId( int prevRemotePublicKeyId ) {
		this.prevRemotePublicKeyId = prevRemotePublicKeyId;
	}

	/**
	 * Returns the current remote dh public key id.
	 * 
	 * @return The current remote dh public key id.
	 */
	public int getCurrentRemotePublicKeyId() {
		return curRemotePublicKeyId;
	}

	/**
	 * Sets the current remote dh public key id.
	 * 
	 * @param currentRemotePublicKeyId The current remote dh public key id to set.
	 */
	public void setCurrentRemotePublicKeyId( int currentRemotePublicKeyId ) {
		this.curRemotePublicKeyId = currentRemotePublicKeyId;
	}

	/**
	 * Returns the current authentication key object.
	 * 
	 * @return The current authentication key object.
	 */
	public JotrAuthKeys getAuthKeys() {
		return authKeys;
	}

	/**
	 * Sets the current authentication key object.
	 * 
	 * @param authKeys The current authentication key object to set.
	 */
	public void setAuthKeys( JotrAuthKeys authKeys ) {
		this.authKeys = authKeys;
	}

	/**
	 * Returns the mac keys to reveal.
	 * 
	 * @return The mac keys to reveal.
	 */
	public byte[] getRevealedMacKeys() {
		return revealedMacKeys;
	}

	/**
	 * Returns the current session keys.
	 * 
	 * @return The current session keys.
	 */
	public JotrSessionKeys[][] getSessKeys() {
		return sessKeys;
	}

	/**
	 * Returns the previous local dh key pair.
	 * 
	 * @return The previous local dh key pair.
	 */
	public JotrDHKeyPair getPreviousLocalDHKeyPair() {
		return prevLocalDHKeyPair;
	}

	/**
	 * Sets the previous local dh key pair.
	 * 
	 * @return The previous local dh key pair to set.
	 */
	public JotrDHKeyPair getCurrentLocalDHKeyPair() {
		return curLocalDHKeyPair;
	}
}